1851G - Vlad and the Mountains - CodeForces Solution

#include <bits/stdc++.h>
using namespace std;
struct DSU{
	vector<int> parent, sz; 
	DSU(int n){
		parent.resize(n + 1);  
		iota(parent.begin(), parent.end(), 0) ;
		sz.assign(n+1, 1) ;  
	}
	void make_set(int v){
		parent[v] = v; 
		sz[v] = 1; 
	}	
	int find_set(int v){
		if(parent[v] == v){
			return v; 
		}
		return parent[v] = find_set(parent[v]) ; 
	}
	bool same(int a, int b){
		return find_set(a) == find_set(b); 
	}
	void merge(int a,int b){
		a = find_set(a) , b = find_set(b);
		if(a != b){
			if(sz[a] < sz[b]) swap(a,b) ; 
			parent[b] = a; 
			sz[a] += sz[b] ; 
		}
	}
};
const int mxN = 2e5; 
vector<vector<int>> all_vertices, check; 
void solve(){
	int n , m ; cin >> n >> m ; 
	vector<int> H(n) ; for(auto &h: H) cin >> h; 
	all_vertices.clear(), check.clear();
	for(int i = 0 ; i < m ; ++i){
		int u ,v ; cin >> u >> v ;
		int f=max(H[u-1],H[v-1]) ; 
		all_vertices.push_back({f,u,v}) ; 
	}
	sort(begin(all_vertices), end(all_vertices)) ; 

	int q; cin >> q; 
	vector<bool> ans(q); 
	for(int i = 0 ; i < q; i ++){
		int a, b, e; cin >> a >> b >> e; 
		// a--,b--;
		//mountain a -> mountain b using energy e
		//H[b] has to be within H[a] + e 
		e += H[a-1];
		check.push_back({e,a,b,i}) ; 
	}
	sort(begin(check), end(check));
	DSU dsu(n) ;
	int j = 0, f = 0 ;  
	for(int i = 0 ; i < check.size() ; i ++){
		while(j < all_vertices.size() && all_vertices[j][0] <= check[i][0]){
			// dsu.make_set(all_vertices[j][1]) , dsu.make_set(all_vertices[j][2]);
			dsu.merge(all_vertices[j][1], all_vertices[j][2]) ;  
			++j;  
		}
		ans[check[i][3]] = dsu.same(check[i][1],check[i][2]) ; 
	}
	//for each query we need to check if for all vertices where max <= e + H[a], b is among them
	for(int i = 0 ; i < q; i ++){
		cout << (ans[i]?"YES\n" : "NO\n") ; 
	}
}
int main(){
	ios_base::sync_with_stdio(false), cin.tie(NULL) ; 
	int t; cin >> t;
	while(t--)solve();
}